package heap;

import java.util.LinkedList;
import java.util.List;
import java.util.PriorityQueue;
import java.util.TreeMap;

/**
 * @program: play-structure
 * @author: baichen
 * 力扣347：给定一个非空的整数数组，返回其中出现频率前 k 高的元素。
 * 如：
 * 输入: nums = [1,1,1,2,2,3], k = 2
 * 输出: [1,2]
 * 参考前面统计频率的类，使用Java提供的优先队列解决，Java提供的默认是最小堆
 **/
public class TopKFrequent1 {
    private class Freq implements Comparable<Freq> {
        int e, freq;     //元素e和频次freq

        public Freq(int e, int freq) {
            this.e = e;
            this.freq = freq;
        }
        //比较器主要是在存入队列时做的比较，这个是最小堆，所以要反过来
        public int compareTo(Freq another) {
            //比较频次大小,要跟自己定义的优先队列区分开，自己定义的是最大堆，Java提供的是最小堆
            if (this.freq < another.freq)
                return -1;
            else if (this.freq > another.freq)
                return 1;
            else        //频次相等
                return 0;
        }
    }

    public List<Integer> topKFrequent(int[] nums, int k) {
        //统计频次
        TreeMap<Integer, Integer> map = new TreeMap<>();
        for (int num : nums) {
            if (map.containsKey(num)) {
                map.put(num, map.get(num) + 1);
            } else
                map.put(num, 1);
        }
        //利用优先队列求出前K个元素,默认是最小堆
        PriorityQueue<Freq> pq = new PriorityQueue<>();
        for (int key : map.keySet()) {
            if (pq.size() < k) {
                pq.add(new Freq(key, map.get(key)));
            } else if (map.get(key) > pq.peek().freq) {
                pq.remove();
                pq.add(new Freq(key, map.get(key)));
            }
        }
        //将前K个元素放入LinkedList
        LinkedList<Integer> result = new LinkedList<>();
        while (!pq.isEmpty()) {
            result.add(pq.remove().e);
        }
        return result;
    }
}
